
#include <vector>
#include "AESSymmetricKey.h"
#include "ECDSAPrivateKey.h"
#include "DSAPrivateKey.h"
#include "RSAPrivateKey.h"
#include "RSAPublicKey.h"
#include "CryptEChain.h"
#include "EncryptWrapper.h"

/** General initialisation function for the encryption wrapper.
 * 
 * @param enckey
 * EncryptionKey to use for encrypting data if present and not NULL.
 * 
 * @param sgnkey
 * SigningKey to use to sign the data stream if present and not NULL.
 * 
 * @param compress
 * A flag to specify an additional compression step if true. */
void EncryptWrapper::initEncryptWrapper(EncryptionKey *enckey, SigningKey *sgnkey, int compress)
{
   if (enckey)
   {
      switch (enckey->getKeyType())
      {
         case Key::KEY_TYPE_PUBLIC_RSA:
            { enckeycopy = new RSAPublicKey(*((RSAPublicKey *)enckey)); break; }
         case Key::KEY_TYPE_SYMM_AES:
            { enckeycopy = new AESSymmetricKey(*((AESSymmetricKey *)enckey)); break; }
         default:
            { errprint("unknown encryption key type"); enckeycopy = NULL; break; }
      }
   } else
   {
      errprint("encrypt wrapper has no key: can't do encryption!");
      enckeycopy = NULL;
   }
   if (sgnkey)
   {
      switch (sgnkey->getKeyType())
      {
         case Key::KEY_TYPE_PRIVATE_RSA:
            { sgnkeycopy = new RSAPrivateKey(*((RSAPrivateKey *)sgnkey)); break; }
         case Key::KEY_TYPE_PRIVATE_DSA:
            { sgnkeycopy = new DSAPrivateKey(*((DSAPrivateKey *)sgnkey)); break; }
         case Key::KEY_TYPE_PRIVATE_ECDSA:
            { sgnkeycopy = new ECDSAPrivateKey(*((ECDSAPrivateKey *)sgnkey)); break; }
         default:
            { errprint("unknown signing key type"); sgnkeycopy = NULL; break; }
      }
   } else
   {
      sgnkeycopy = NULL;
   }
   docompress = (compress != 0);
   trusted = 0;
   apr_status_t status;
   status = apr_pool_create(&mp, NULL);
   CheckAPRError(status);
}

/** Constructs an encryption process which encrypts with the given
 * EncryptionKey.
 * 
 * @param enckey
 * EncryptionKey to use for encrypting data if present and not NULL. */
EncryptWrapper::EncryptWrapper(EncryptionKey *enckey)
{
   initEncryptWrapper(enckey, NULL, FALSE);
}

/** Constructs an encryption process which encrypts the data with the given
 * EncryptionKey and subsequently signs the result with the given SigningKey.
 * 
 * @param enckey
 * EncryptionKey to use for encrypting data if present and not NULL.
 * 
 * @param sgnkey
 * SigningKey to use to sign the data stream if present and not NULL. */
EncryptWrapper::EncryptWrapper(EncryptionKey *enckey, SigningKey *sgnkey)
{
   initEncryptWrapper(enckey, sgnkey, FALSE);
}

/** Constructs an encryption process which can optionally apply a compression
 * step, before encrypting with the given EncryptionKey.
 * 
 * @param enckey
 * EncryptionKey to use for encrypting data if present and not NULL.
 * 
 * @param docompress
 * A flag to specify an additional compression step if true. */
EncryptWrapper::EncryptWrapper(EncryptionKey *enckey, int docompress)
{
   initEncryptWrapper(enckey, NULL, docompress);
}

/** Constructs an encryption process which can optionally apply a compression
 * step, before encrypting with the given EncryptionKey and signing the result
 * with the given SigningKey.
 * 
 * @param enckey
 * EncryptionKey to use for encrypting data if present and not NULL.
 * 
 * @param sgnkey
 * SigningKey to use to sign the data stream if present and not NULL.
 * 
 * @param docompress
 * A flag to specify an additional compression step if true. */
EncryptWrapper::EncryptWrapper(EncryptionKey *enckey, SigningKey *sgnkey, int docompress)
{
   initEncryptWrapper(enckey, sgnkey, docompress);
}

/** Destructor for the encryption wrapper. */
EncryptWrapper::~EncryptWrapper()
{
   if (enckeycopy) delete enckeycopy;
   if (sgnkeycopy) delete sgnkeycopy;
   apr_pool_destroy(mp);
}

/** Apply the encryption process to a buffer of data, returning a vector of
 * encoded data.
 * 
 * @param buffer
 * Buffer for uncoded data input.
 * 
 * @param length
 * Length of buffer for data input.
 * 
 * @return
 * An std::vector representing a variable-sized buffer into which the encoded
 * data is put. The caller is responsible for deleting the vector instance. */
std::vector<char> *EncryptWrapper::encryptBuffer(const char *buffer, size_t length)
{
   int smallblock = (ENV_READ_BUFFER > (length/2)) ? (length/2) : ENV_READ_BUFFER;
   size_t nlength = length;
   APRBucketBrigade *inputbb = new APRBucketBrigade(mp);
   if (inputbb->bufferToBucketBrigade(buffer, &length, smallblock) != 0)
   {
      errprint("failed to push the buffer into a bucket");
      delete inputbb;
      return NULL;
   }
   if (nlength != length)
   {
      errprint("bucket generation failed, pushed %d bytes to bucket, wanted %d", nlength, length);
      delete inputbb;
      return NULL;
   }
   APRBucketBrigade *outputbb = encryptToBucket(inputbb);
   delete inputbb;
   char *outputbuffer = NULL;
   size_t ilength = 0;
   std::vector<char> *newstring = NULL;
   if (outputbb->bucketBrigadeToBuffer(&outputbuffer, &ilength) == 0)
   {
      // outputbuffer cannot be NULL!
      newstring = new std::vector<char>();
      newstring->insert(newstring->end(), outputbuffer, &(outputbuffer[ilength]));
   } else
   {
      errprint("failed to put bucket into buffer");
   }
   if (outputbuffer) free(outputbuffer);
   delete outputbb;
   return newstring;
}

/** Apply the encryption process to a pipe of data in the form of an
 * APRBucketBrigade, returning a pointer to a newly created APRBucketBrigade
 * containing the encoded data.
 * 
 * @param bbrgde
 * The bucket brigade input of uncoded data.
 * 
 * @return
 * New bucket brigade representing the complete encoded data stream. The
 * caller is responsible for deleting the APRBucketBrigade instance. */
APRBucketBrigade *EncryptWrapper::encryptToBucket(APRBucketBrigade *bbrgde)
{
   CryptEChain *crypty = new CryptEChain();
   crypty->buildEncryptionChain2(enckeycopy, sgnkeycopy, docompress, bbrgde->realLength());
   APRBucketBrigade *outputbrgde = new APRBucketBrigade(mp);
   int prociter = 0;
   while (prociter == 0)
      prociter = crypty->processIteration(bbrgde, outputbrgde);
   delete bbrgde;
   trusted = crypty->isTrusted();
   delete crypty;
   if (prociter == -1)
   {
      delete outputbrgde;
      return NULL;
   }
   return outputbrgde;
}
